# -*- coding: utf-8 -*-
from ultralytics.models import YOLO
import mlflow
import waveletai
import torch
from waveletai.constants import FeatureType
import os
from mlflow.utils.file_utils import TempDir
import shutil
import click
import sys
from pathlib import Path

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

# os.environ["WAVELETAI_API_URL"] = "http://192.168.2.89/api"


@click.command()
@click.option("--epochs", "-ep", type=int, default=1)
@click.option("--batch_size", "-bs", type=int, default=64)
@click.option("--weights_file", "-wf", type=str, default='yolov8n-seg.pt')
@click.option("--training_data", "-td", type=str, default="", help="数据集id")
def enter(epochs, batch_size, weights_file, training_data):
    pwd = os.path.dirname(os.path.abspath(__file__))
    sys.path.insert(0, pwd)
    os.chdir(pwd)

    # with mlflow.start_run(nested=True):
    waveletai.init()
    if weights_file.endswith(".pt"):
        weights_file = weights_file
    else:
        pretrain_model = weights_file
        if (pretrain_model):
            print("pretrain_model processing..........")
            #waveletai.set_model("1a7cdc805ad641c0854911b788ea8247")
            waveletai.download_model_version_asset(pretrain_model, 'artifacts/best.pt', "./pretrain_model")
            assert os.path.exists("pretrain_model/best.pt"), "best.pt not exists"
            weights_file = 'pretrain_model/best.pt'
    waveletai.download_dataset_artifacts(training_data, './dataset', unzip=True,
                                            feature_type=FeatureType.YOLOV5.value)
    # 数据处理
    with TempDir() as tmp:
        print("dataset processing..........")
        tmp_path = tmp.path()
        output_directory = tmp_path + '/'
        print("output_directory: ", output_directory)
        waveletai.download_dataset_artifacts(training_data,
                                                    output_directory + '/dataset', unzip=True,
                                                    feature_type=FeatureType.YOLOV5.value)
        raw_train = output_directory + 'dataset/train'
        raw_valid = output_directory + 'dataset/valid'

        
        if not os.path.isdir(output_directory +"data/train/images"):
            os.makedirs(output_directory +"data/train/images")
        if not os.path.isdir(output_directory +"data/train/labels"):
            os.makedirs(output_directory +"data/train/labels")
        if not os.path.isdir(output_directory +"data/valid/images"):
            os.makedirs(output_directory +"data/valid/images")
        if not os.path.isdir(output_directory +"data/valid/labels"):
            os.makedirs(output_directory +"data/valid/labels")
        if not os.path.isdir(output_directory +"data/test/images"):
            os.makedirs(output_directory +"data/test/images")
        if not os.path.isdir(output_directory +"data/test/labels"):
            os.makedirs(output_directory +"data/test/labels")


        files = os.listdir(output_directory + 'dataset')
        print(files)
          
        list_train = os.listdir(raw_train)
        print("list_train: ", len(list_train))
        for i in range(len(list_train)):
            if list_train[i][-4:] == '.jpg':
                old_name = raw_train + '/' + str(list_train[i])
                new_name = output_directory + 'data/train/images/' + str(list_train[i])
                shutil.copyfile(old_name, new_name)
            elif list_train[i][-4:] == '.txt':
                old_name = raw_train + '/' + str(list_train[i])
                new_name = output_directory + 'data/train/labels/' + str(list_train[i])
                shutil.copyfile(old_name, new_name)
        print('完成 train - 复制到images/复制到labels')
        

        if not os.path.isdir(raw_valid):
            list_valid = os.listdir(raw_train)
            raw_valid = raw_train
        else:
            list_valid = os.listdir(raw_valid)
        print("list_valid: ", len(list_valid))

        for i in range(len(list_valid)):
            if list_valid[i][-4:] == '.jpg':
                old_name = raw_valid + '/' + str(list_valid[i])
                new_name = output_directory + 'data/valid/images/' + str(list_valid[i])
                shutil.copyfile(old_name, new_name)
            elif list_valid[i][-4:] == '.txt':
                old_name = raw_valid + '/' + str(list_valid[i])
                new_name = output_directory + 'data/valid/labels/' + str(list_valid[i])
                shutil.copyfile(old_name, new_name)

        for i in range(len(list_train)):
            if list_train[i][-4:] == '.jpg':
                old_name = raw_train + '/' + str(list_train[i])
                new_name = output_directory + 'data/valid/images/' + str(list_train[i])
                shutil.copyfile(old_name, new_name)
            elif list_train[i][-4:] == '.txt':
                old_name = raw_train + '/' + str(list_train[i])
                new_name = output_directory + 'data/valid/labels/' + str(list_train[i])
                shutil.copyfile(old_name, new_name)
    
        print('完成 valid - 复制到images/复制到labels')
        # 数据处理结束
        print("创建对象")
        # 保存参数
        waveletai.log_param("batch_size_value", batch_size)
        waveletai.log_param("epochs_value", epochs)
        # 训练模型
        model = YOLO(weights_file)
        device="cuda:0" if torch.cuda.is_available() else "cpu"
        model.train(data=output_directory + '/dataset/data.yaml', epochs=epochs, imgsz=640, device=device, project=str(pwd))


if __name__ == "__main__":
    # enter(epochs=3, batch_size=10, weights_file='yolov8n.pt', training_data='2b525bc4491b4f038a165e91770c034e')
    enter()

